#ifndef _DatasetIO_dataset_io_nclt_H_
#define _DatasetIO_dataset_io_nclt_H_

#include "dataset_io.h"
#include "dataset_io_utils.h"

namespace DatasetIO {

class NcltDataset : public Dataset {
    std::string path;
    std::string sequence_id;

    // Image
    std::string image_base_path;
    size_t num_cams;
    std::vector<double> image_timestamps;
    std::unordered_map<double, std::string> image_path;
    size_t cur_image_id = 0;

    // Lidar
    std::vector<double> lidar_timestamps;
    std::unordered_map<double, std::string> lidar_path;
    size_t cur_lidar_id = 0;

    // Imu
    std::vector<double> imu_timestamps;
    std::vector<AccelData::Ptr> accel_data;
    std::vector<GyroData::Ptr> gyro_data;
    size_t cur_imu_id = 0;

    // Wheel
    std::vector<double> wheel_timestamps;
    std::vector<WheelData::Ptr> wheel_data;
    size_t cur_wheel_id = 0;

    // Gnss
    std::vector<double> gnss_timestamps;
    std::vector<GnssData::Ptr> gnss_data;
    size_t cur_gnss_id = 0;

    // GroundTrue
    std::vector<double> gt_timestamps;  // ordered gt timestamps
    std::vector<PoseData::Ptr> gt_pose_data;

    const std::vector<std::string> folder = {"/Cam5/", "/Cam4/", "/Cam3/",
                                             "/Cam2/", "/Cam1/"};

public:
    NcltDataset() : Dataset() { type = DatasetType::kNclt; }
    ~NcltDataset(){};

    std::vector<double> &get_image_timestamps() { return image_timestamps; }

    const std::vector<AccelData::Ptr> &get_accel_data() const {
        return accel_data;
    }
    const std::vector<GyroData::Ptr> &get_gyro_data() const {
        return gyro_data;
    }

    std::vector<std::vector<Measurement::Ptr>> get_nearest_data() {
        std::vector<std::pair<double, int>> measurement_timestamp;

        if (cur_image_id < image_timestamps.size())
            measurement_timestamp.push_back(
                std::pair<double, int>(image_timestamps[cur_image_id], 0));
        if (cur_lidar_id < lidar_timestamps.size())
            measurement_timestamp.push_back(
                std::pair<double, int>(lidar_timestamps[cur_lidar_id], 1));
        if (cur_imu_id < imu_timestamps.size())
            measurement_timestamp.push_back(
                std::pair<double, int>(imu_timestamps[cur_imu_id], 2));
        if (cur_wheel_id < wheel_timestamps.size())
            measurement_timestamp.push_back(
                std::pair<double, int>(wheel_timestamps[cur_wheel_id], 3));
        if (cur_gnss_id < gnss_timestamps.size())
            measurement_timestamp.push_back(
                std::pair<double, int>(gnss_timestamps[cur_gnss_id], 4));

        std::sort(measurement_timestamp.begin(), measurement_timestamp.end(),
                  timestamp_comp);

        std::vector<std::vector<Measurement::Ptr>> all_measurements;

        for (size_t i = 0; i < measurement_timestamp.size(); i++) {
            if (std::abs(measurement_timestamp[i].first -
                         measurement_timestamp[0].first) < 0.0015) {
                all_measurements.push_back(
                    get_next_data(measurement_timestamp[i].second));
            }
        }

        return all_measurements;
    }

    std::vector<MonoImageData::Ptr> get_image_data(double t_ns) {
        std::vector<MonoImageData::Ptr> res(num_cams);

        if (image_path.find(t_ns) != image_path.end()) {
            for (size_t i = 0; i < num_cams; i++) {
                std::string full_image_path =
                    image_base_path + folder[i] + image_path[t_ns] + ".tiff";

                if (fs::exists(full_image_path)) {
                    cv::Mat img;
                    if (!(img =
                              cv::imread(full_image_path, cv::IMREAD_UNCHANGED))
                             .empty()) {
                        res[i] = std::make_shared<MonoImageData>();
                        res[i]->timestamp = t_ns;
                        res[i]->data = img.clone();
                    }
                }
            }
        }

        return res;
    }

    MonoImageData::Ptr get_next_mono_image_data() {
        MonoImageData::Ptr res;
        cv::Mat image;

        double t_s = image_timestamps[cur_image_id++];

        std::string full_image_path =
            image_base_path + folder[0] + image_path[t_s];
        if (!(image = cv::imread(full_image_path, cv::IMREAD_UNCHANGED))
                 .empty()) {
            res = std::make_shared<MonoImageData>();
            res->timestamp = t_s;
            res->data = image.clone();
        }

        return res;
    }

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    friend class NcltIO;

protected:
    std::vector<Measurement::Ptr> get_next_data(int index) {
        std::vector<Measurement::Ptr> next_data;

        if (index == 0) {  // Mono image
            double ts = image_timestamps[cur_image_id++];
            auto image_data_temp = get_image_data(ts);
            for (int i = 0; i < image_data_temp.size(); i++) {
                next_data.push_back(Measurement::Ptr(image_data_temp[i]));
            }
        } else if (index == 1) {  // Lidar
            double ts = lidar_timestamps[cur_lidar_id++];
            // auto lidar_data_temp = get_image_data(ts);
            // for (int i = 0; i < lidar_data_temp.size(); i++) {
            //     next_data.push_back(Measurement::Ptr(lidar_data_temp[i]));
            // }
        } else if (index == 2) {  // Imu
            Measurement::Ptr d1(accel_data[cur_imu_id]);
            Measurement::Ptr d2(gyro_data[cur_imu_id++]);
            next_data.push_back(d1);
            next_data.push_back(d2);
        } else if (index == 3) {  // Wheel
            Measurement::Ptr d1(wheel_data[cur_wheel_id++]);
            next_data.push_back(d1);
        } else if (index == 4) {  // GNSS
            Measurement::Ptr d1(gnss_data[cur_gnss_id++]);
            next_data.push_back(d1);
        }

        return next_data;
    }
};

class NcltIO : public DatasetIoInterface {
public:
    NcltIO() {}

    void read(const std::string &path) {
        if (!fs::exists(path))
            LOG(WARNING) << "No dataset found in " << path;

        data.reset(new NcltDataset);

        data->num_cams = 5;
        data->path = path;
        data->cur_image_id = 0;

        auto data_info_node = YAML::LoadFile(path + "/test_setting.yaml");
        std::string sequence_id =
            data_info_node["sequence_id"].as<std::string>();
        data->sequence_id = sequence_id;

        data->image_base_path =
            path + "/" + sequence_id + "/" + sequence_id + "/lb3/";
        std::string imu_path = path + "/" + sequence_id + "/" + sequence_id +
                               "_sen/" + sequence_id + "/ms25.csv";
        std::string wheel_path = path + "/" + sequence_id + "/" + sequence_id +
                                 "_sen/" + sequence_id + "/wheels.csv";
        std::string gnss_path = path + "/" + sequence_id + "/" + sequence_id +
                                "_sen/" + sequence_id + "/gps_rtk.csv";
        std::string gt_path =
            path + "/" + sequence_id + "/groundtruth_" + sequence_id + ".csv";

        read_image_timestamps(path + "/rgb/rgb_LB_nclt_" + sequence_id +
                              ".txt");
        read_imu(imu_path);
        read_wheel(wheel_path);
        read_gnss(gnss_path);
        read_gt_data_pose(gt_path);
    }

    void reset() { data.reset(); }

    Dataset::Ptr get_data() { return data; }

private:
    void read_image_timestamps(const std::string &path) {
        std::string strPathTimeFile = path;

        std::ifstream fTimes;
        fTimes.open(strPathTimeFile.c_str());
        if (!fTimes.is_open()) {
            LOG(ERROR) << "Can't open " << strPathTimeFile;
            exit(EXIT_FAILURE);
        }

        int nTimes = 0;
        while (!fTimes.eof()) {
            std::string s;
            std::getline(fTimes, s);

            if (s[0] == '#' || s == "")
                continue;

            if (!s.empty()) {
                std::stringstream ss;
                ss << s;
                double t;
                ss >> t;
                double t_s = t / 1e6;
                data->image_timestamps.emplace_back(t_s);
                data->image_path[t_s] = s;

                nTimes++;
            }
        }

        LOG(INFO) << "Reading over " << strPathTimeFile;
    }

    void read_imu(const std::string &path) {
        std::string strImuPath = path;

        std::ifstream fImu;
        fImu.open(strImuPath.c_str());
        if (!fImu.is_open()) {
            LOG(ERROR) << "Can't open " << strImuPath;
            exit(EXIT_FAILURE);
        }

        while (!fImu.eof()) {
            std::string s;
            std::getline(fImu, s);
            if (s[0] == '#' || s == "")
                continue;

            if (!s.empty()) {
                std::string item;
                size_t pos = 0;
                double item_data[10];
                int count = 0;
                while ((pos = s.find(',')) != std::string::npos) {
                    item = s.substr(0, pos);
                    item_data[count++] = stod(item);
                    s.erase(0, pos + 1);
                }
                item = s.substr(0, pos);
                item_data[9] = stod(item);

                double ts = item_data[0] / 1e6;
                data->imu_timestamps.push_back(ts);
                AccelData::Ptr acc(new AccelData);
                acc->timestamp = ts;
                acc->data =
                    Eigen::Vector3d(item_data[4], item_data[5], item_data[6]);
                data->accel_data.push_back(acc);
                GyroData::Ptr gyr(new GyroData);
                gyr->timestamp = ts;
                gyr->data =
                    Eigen::Vector3d(item_data[7], item_data[8], item_data[9]);
                data->gyro_data.push_back(gyr);
            }
        }

        LOG(INFO) << "Reading over " << strImuPath;
    }

    void read_wheel(const std::string &path) {
        std::string strWheelPath = path;

        std::ifstream fWheel;
        fWheel.open(strWheelPath.c_str());
        if (!fWheel.is_open()) {
            LOG(ERROR) << "Can't open " << strWheelPath;
            exit(EXIT_FAILURE);
        }

        while (!fWheel.eof()) {
            std::string s;
            std::getline(fWheel, s);
            if (s[0] == '#' || s == "")
                continue;

            if (!s.empty()) {
                std::string item;
                size_t pos = 0;
                double item_data[3];
                int count = 0;
                while ((pos = s.find(',')) != std::string::npos) {
                    item = s.substr(0, pos);
                    item_data[count++] = stod(item);
                    s.erase(0, pos + 1);
                }
                item = s.substr(0, pos);
                item_data[2] = stod(item);

                double ts = item_data[0] / 1e6;
                data->wheel_timestamps.push_back(ts);
                WheelData::Ptr wheel(new WheelData);
                wheel->timestamp = ts;
                wheel->data = Eigen::Vector2d(item_data[1], item_data[2]);
                data->wheel_data.push_back(wheel);
            }
        }

        LOG(INFO) << "Reading over " << strWheelPath;
    }

    void read_gnss(const std::string &path) {
        std::string strGnssPath = path;

        std::ifstream fGnss;
        fGnss.open(strGnssPath.c_str());
        if (!fGnss.is_open()) {
            LOG(ERROR) << "Can't open " << strGnssPath;
            exit(EXIT_FAILURE);
        }

        while (!fGnss.eof()) {
            std::string s;
            getline(fGnss, s);
            if (s[0] == '#' || s == "")
                continue;

            if (!s.empty()) {
                std::string item;
                size_t pos = 0;
                double item_data[8];
                int count = 0;
                while ((pos = s.find(',')) != std::string::npos) {
                    item = s.substr(0, pos);
                    item_data[count++] = stod(item);
                    s.erase(0, pos + 1);
                }
                item = s.substr(0, pos);
                item_data[7] = stod(item);

                if (item_data[1] < 3)
                    continue;

                if (std::isnan(item_data[3]) || std::isnan(item_data[4]) ||
                    std::isnan(item_data[5]) || std::isnan(item_data[6]) ||
                    std::isnan(item_data[7]))
                    continue;

                double ts = item_data[0] / 1e6;
                data->gnss_timestamps.push_back(ts);
                GnssData::Ptr gnss(new GnssData);
                gnss->timestamp = ts;
                gnss->data =
                    Eigen::Vector3d(item_data[3] * 180. / M_PI,
                                    item_data[4] * 180. / M_PI, item_data[5]);
                data->gnss_data.push_back(gnss);
            }
        }

        LOG(INFO) << "Reading over " << strGnssPath;
    }

    void read_gt_data_pose(const std::string &path) {
        std::string strGtPath = path;

        std::ifstream fGt;
        fGt.open(strGtPath.c_str());
        if (!fGt.is_open()) {
            LOG(ERROR) << "Can't open " << strGtPath;
            exit(EXIT_FAILURE);
        }

        while (!fGt.eof()) {
            std::string s;
            std::getline(fGt, s);
            if (s[0] == '#' || s == "")
                continue;

            if (!s.empty()) {
                std::string item;
                size_t pos = 0;
                double item_data[7];
                int count = 0;
                while ((pos = s.find(',')) != std::string::npos) {
                    item = s.substr(0, pos);
                    item_data[count++] = stod(item);
                    s.erase(0, pos + 1);
                }
                item = s.substr(0, pos);
                item_data[6] = stod(item);

                double ts = item_data[0] / 1e6;
                data->gt_timestamps.push_back(ts);

                Eigen::Matrix4d pose = Eigen::Matrix4d::Identity();
                // pose.block<3,3>(0,0) = Ei;
                pose.block<3, 1>(0, 3) =
                    Eigen::Vector3d(item_data[2], item_data[1], -item_data[3]);

                PoseData::Ptr pose_data(new PoseData);
                pose_data->timestamp = ts;
                pose_data->data = pose;

                data->gt_pose_data.push_back(pose_data);
            }
        }

        LOG(INFO) << "Reading over " << strGtPath;
    }

    std::shared_ptr<NcltDataset> data;
};

}  // namespace DatasetIO

#endif